package com.rickpan.service;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.rickpan.dto.openrouter.OpenRouterRequest;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.*;
import org.springframework.stereotype.Component;
import org.springframework.web.client.RestTemplate;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.nio.charset.StandardCharsets;
import java.util.function.Consumer;

/**
 * OpenRouter API客户端
 */
@Component
public class OpenRouterClient {

    private static final Logger log = LoggerFactory.getLogger(OpenRouterClient.class);
    private static final String OPENROUTER_API_URL = "https://openrouter.ai/api/v1/chat/completions";
    
    @Autowired
    private RestTemplate restTemplate;
    
    @Autowired
    private ObjectMapper objectMapper;

    /**
     * 发送流式请求到OpenRouter
     * @param apiKey API密钥
     * @param request 请求对象
     * @param chunkHandler 处理流式数据的回调函数
     */
    public void sendStreamRequest(String apiKey, OpenRouterRequest request, Consumer<String> chunkHandler) {
        HttpHeaders headers = new HttpHeaders();
        headers.set("Authorization", "Bearer " + apiKey);
        headers.set("Content-Type", "application/json");
        headers.set("Accept", "text/event-stream");

        try {
            // 使用RestTemplate的execute方法处理流式响应
            restTemplate.execute(OPENROUTER_API_URL, HttpMethod.POST, 
                requestCallback -> {
                    try {
                        requestCallback.getHeaders().putAll(headers);
                        objectMapper.writeValue(requestCallback.getBody(), request);
                    } catch (IOException e) {
                        log.error("Error writing request body: {}", e.getMessage(), e);
                        throw new RuntimeException("Request writing error", e);
                    }
                },
                responseExtractor -> {
                    try (BufferedReader reader = new BufferedReader(
                            new InputStreamReader(responseExtractor.getBody(), StandardCharsets.UTF_8))) {
                        String line;
                        while ((line = reader.readLine()) != null) {
                            if (!line.trim().isEmpty()) {
                                chunkHandler.accept(line);
                            }
                        }
                    } catch (IOException e) {
                        log.error("Error reading stream response: {}", e.getMessage(), e);
                        throw new RuntimeException("Stream reading error", e);
                    }
                    return null;
                });
        } catch (Exception e) {
            log.error("Error in stream request: {}", e.getMessage(), e);

            // 特殊处理429错误（限额超出）
            if (e.getMessage() != null && e.getMessage().contains("429")) {
                throw new RuntimeException("API调用限额已用完，请稍后重试或充值积分", e);
            }

            throw new RuntimeException("OpenRouter API call failed: " + e.getMessage(), e);
        }
    }

    /**
     * 测试API Key是否有效
     * @param apiKey API密钥
     * @return 是否有效
     */
    public boolean testApiKey(String apiKey) {
        try {
            // 创建简单的测试请求
            String testRequestJson = """
                {
                    "model": "deepseek/deepseek-chat-v3-0324:free",
                    "messages": [
                        {
                            "role": "user",
                            "content": "Hello"
                        }
                    ],
                    "max_tokens": 10
                }
                """;

            HttpHeaders headers = new HttpHeaders();
            headers.set("Authorization", "Bearer " + apiKey);
            headers.set("Content-Type", "application/json");

            HttpEntity<String> entity = new HttpEntity<>(testRequestJson, headers);
            ResponseEntity<String> response = restTemplate.exchange(
                    OPENROUTER_API_URL, HttpMethod.POST, entity, String.class);

            return response.getStatusCode().is2xxSuccessful();
        } catch (Exception e) {
            log.warn("API Key test failed: {}", e.getMessage());
            return false;
        }
    }

    /**
     * 解析流式响应中的内容
     * @param chunk 流式数据块
     * @return 解析出的内容，如果无内容返回null
     */
    public String parseStreamChunk(String chunk) {
        try {
            if (chunk.startsWith("data: ")) {
                String data = chunk.substring(6).trim();
                if ("[DONE]".equals(data) || data.isEmpty()) {
                    return null;
                }
                
                // 解析JSON并提取content
                JsonNode jsonNode = objectMapper.readTree(data);
                JsonNode choices = jsonNode.get("choices");
                if (choices != null && choices.isArray() && choices.size() > 0) {
                    JsonNode delta = choices.get(0).get("delta");
                    if (delta != null && delta.has("content")) {
                        return delta.get("content").asText();
                    }
                }
            }
        } catch (Exception e) {
            log.warn("Error parsing stream chunk: {}", e.getMessage());
        }
        return null;
    }
}
